from manim import *
import numpy as np

class ImprovedImageRecognition(Scene):
    def construct(self):
        # 设置背景颜色
        self.camera.background_color = "#1e1e1e"
        
        # 标题
        title = Text("大模型图像识别原理", font_size=36, color=WHITE)
        self.play(Write(title))
        self.wait(1)
        self.play(FadeOut(title))
        
        # 分步骤展示
        self.show_input_image()
        self.show_preprocessing()
        self.show_convolution()
        self.show_feature_extraction()
        self.show_classification()
    
    def show_input_image(self):
        # 清除可能存在的残留元素
        self.clear()
        
        # 输入图像
        image = Rectangle(height=2, width=3, fill_color=BLUE, fill_opacity=1, 
                         stroke_color=WHITE, stroke_width=2)
        image_label = Text("输入图像", font_size=24, color=WHITE).next_to(image, DOWN)
        
        self.play(FadeIn(image), Write(image_label))
        self.wait(1)
        
        self.image_group = VGroup(image, image_label)
        self.play(self.image_group.animate.to_edge(UP))
        self.wait(0.5)
    
    def show_preprocessing(self):
        # 完全清除之前的场景
        self.play(FadeOut(self.image_group))
        
        # 创建像素网格
        pixel_grid = self.create_pixel_grid(6, 8)
        pixel_grid.scale(0.8).to_edge(UP)
        
        # 预处理标签
        preprocess_label = Text("1. 图像预处理", font_size=24, color=WHITE)
        detail_label = Text("转换为像素值矩阵", font_size=20, color=GRAY_A).next_to(preprocess_label, DOWN)
        labels = VGroup(preprocess_label, detail_label).next_to(pixel_grid, DOWN, buff=0.5)
        
        # 动画展示
        self.play(
            FadeIn(pixel_grid, shift=UP),
            Write(labels),
            run_time=2
        )
        self.wait(1)
        
        # 转换为张量
        tensor_cube = self.create_tensor_cube()
        tensor_cube.scale(0.7).next_to(pixel_grid, RIGHT, buff=1)
        tensor_label = Text("3D张量 (H×W×C)", font_size=20, color=WHITE).next_to(tensor_cube, DOWN)
        
        self.play(
            TransformFromCopy(pixel_grid, tensor_cube),
            Write(tensor_label),
            run_time=2
        )
        self.wait(1)
        
        # 保存预处理组
        self.preprocess_group = VGroup(pixel_grid, labels, tensor_cube, tensor_label)
        self.wait(1)
    
    def show_convolution(self):
        # 完全清除预处理内容
        self.play(FadeOut(self.preprocess_group))
        
        # 创建新的输入矩阵（确保位置准确）
        input_matrix = self.create_pixel_grid(6, 8, color=BLUE)
        input_matrix.scale(0.6).to_edge(LEFT).shift(UP*0.5)
        input_label = Text("输入矩阵", font_size=20, color=WHITE).next_to(input_matrix, DOWN)
        
        # 创建卷积核
        kernel = self.create_conv_kernel()
        kernel.scale(0.7).to_edge(RIGHT)
        kernel_label = Text("卷积核", font_size=20, color=WHITE).next_to(kernel, DOWN)
        
        # 标题
        conv_title = Text("2. 卷积操作", font_size=28, color=YELLOW).to_edge(UP)
        detail = Text("滑动窗口提取局部特征", font_size=20, color=GRAY_A).next_to(conv_title, DOWN)
        
        # 动画展示
        self.play(
            Write(conv_title),
            Write(detail),
            run_time=1
        )
        self.play(
            FadeIn(input_matrix),
            Write(input_label),
            FadeIn(kernel),
            Write(kernel_label),
            run_time=1
        )
        self.wait(0.5)
        
        # 精确的卷积动画
        self.animate_precise_convolution(input_matrix, kernel)
        
        # 展示多个卷积核
        kernels = VGroup(*[
            self.create_conv_kernel(color=color).scale(0.6) 
            for color in [RED, GREEN, YELLOW, PURPLE]
        ])
        kernels.arrange(RIGHT, buff=0.3)
        kernels.next_to(detail, DOWN, buff=1)
        multi_label = Text("多个卷积核提取不同特征", font_size=20, color=WHITE).next_to(kernels, DOWN)
        
        self.play(
            LaggedStart(*[GrowFromCenter(k) for k in kernels], lag_ratio=0.2),
            Write(multi_label),
            run_time=2
        )
        self.wait(2)
        
        # 保存当前内容
        self.conv_group = VGroup(
            conv_title, detail, 
            input_matrix, input_label,
            kernel, kernel_label,
            kernels, multi_label
        )
    
    def animate_precise_convolution(self, input_matrix, kernel):
        """精确的卷积动画，修正了位置计算问题"""
        # 创建卷积核副本用于动画
        kernel_copy = kernel.copy()
        self.add(kernel_copy)
        
        # 计算输入矩阵的网格参数
        rows = 6
        cols = 8
        square_size = 0.3 * 0.6  # 原始大小0.3，缩放0.6
        
        # 高亮显示卷积核覆盖区域
        highlight = SurroundingRectangle(
            VGroup(*input_matrix[:3]), 
            color=YELLOW, buff=0, stroke_width=2
        )
        self.add(highlight)
        
        # 初始位置
        start_pos = input_matrix[0].get_center() + RIGHT*(square_size) + DOWN*(square_size)
        kernel_copy.move_to(start_pos)
        
        self.play(
            FadeIn(kernel_copy),
            run_time=0.5
        )
        self.wait(0.5)
        
        # 卷积核滑动动画
        for i in range(rows-2):
            for j in range(cols-2):
                # 计算精确的目标位置
                target_pos = input_matrix[i*cols + j].get_center() + \
                            RIGHT*(square_size) + DOWN*(square_size)
                
                # 更新高亮区域
                new_highlight = SurroundingRectangle(
                    VGroup(*[
                        input_matrix[(i+k)*cols + j + l] 
                        for k in range(3) for l in range(3)
                    ]),
                    color=YELLOW, buff=0, stroke_width=2
                )
                
                self.play(
                    kernel_copy.animate.move_to(target_pos),
                    Transform(highlight, new_highlight),
                    run_time=0.5
                )
        
        # 清除动画对象
        self.play(
            FadeOut(highlight),
            FadeOut(kernel_copy),
            run_time=0.5
        )
    
    def show_feature_extraction(self):
        # 完全清除卷积内容
        self.play(FadeOut(self.conv_group))
        
        # 创建特征图
        feature_maps = self.create_feature_maps()
        feature_maps.scale(0.7).center()
        
        # 标题和说明
        title = Text("3. 特征提取", font_size=28, color=YELLOW).to_edge(UP)
        detail = Text("卷积后得到的特征图", font_size=20, color=GRAY_A).next_to(title, DOWN)
        
        # 定位
        feature_label = Text("特征图", font_size=20, color=WHITE).next_to(feature_maps, DOWN)
        
        # 动画展示
        self.play(
            Write(title),
            Write(detail),
            run_time=1
        )
        self.play(
            FadeIn(feature_maps),
            Write(feature_label),
            run_time=2
        )
        self.wait(2)
        
        # 保存当前内容
        self.feature_group = VGroup(title, detail, feature_maps, feature_label)
    
    def show_classification(self):
        # 完全清除特征图
        self.play(FadeOut(self.feature_group))
        
        # 创建全连接层
        fc_layers = self.create_fc_layers()
        fc_layers.scale(0.8).to_edge(LEFT)
        
        # 创建分类结果
        output = Text("猫: 0.85\n狗: 0.10\n鸟: 0.05", font_size=24, color=WHITE)
        output.next_to(fc_layers, RIGHT, buff=1)
        output_label = Text("分类结果", font_size=24, color=YELLOW).next_to(output, UP)
        
        # 标题
        title = Text("4. 分类预测", font_size=28, color=YELLOW).to_edge(UP)
        detail = Text("全连接层 + Softmax", font_size=20, color=GRAY_A).next_to(title, DOWN)
        
        # 定位
        fc_label = Text("全连接层", font_size=20, color=WHITE).next_to(fc_layers, UP)
        
        # 动画展示
        self.play(
            Write(title),
            Write(detail),
            run_time=1
        )
        self.play(
            FadeIn(fc_layers),
            Write(fc_label),
            FadeIn(output),
            Write(output_label),
            run_time=2
        )
        
        # 添加连接动画
        self.play(
            Circumscribe(output, color=GREEN, run_time=2),
            Flash(output, color=YELLOW, run_time=2)
        )
        self.wait(3)
    
    # 辅助函数保持不变
    def create_pixel_grid(self, rows, cols, color=BLUE):
        grid = VGroup()
        for i in range(rows):
            for j in range(cols):
                pixel = Square(side_length=0.3, fill_color=color,
                             fill_opacity=np.random.uniform(0.3, 1),
                             stroke_color=WHITE, stroke_width=1)
                pixel.move_to([j*0.3 - (cols*0.3)/2, -i*0.3 + (rows*0.3)/2, 0])
                grid.add(pixel)
        return grid
    
    def create_tensor_cube(self):
        cube = Cube(side_length=1.5, fill_color=BLUE_E, fill_opacity=0.8, 
                   stroke_color=WHITE, stroke_width=2)
        
        for i in np.linspace(-0.6, 0.6, 3):
            line = Line([i, -0.75, -0.75], [i, 0.75, -0.75], color=WHITE, stroke_width=1)
            cube.add(line)
            line = Line([-0.75, i, -0.75], [0.75, i, -0.75], color=WHITE, stroke_width=1)
            cube.add(line)
        
        return cube
    
    def create_conv_kernel(self, color=RED):
        kernel = VGroup()
        for i in range(3):
            for j in range(3):
                val = np.random.uniform(-1, 1)
                pixel = Square(side_length=0.35, fill_color=color,
                              fill_opacity=abs(val)*0.5 + 0.3,
                              stroke_color=WHITE, stroke_width=1.5)
                pixel.move_to([j*0.35 - 0.35, -i*0.35 + 0.35, 0])
                text = Text(f"{val:.1f}", font_size=14, color=WHITE).move_to(pixel)
                kernel.add(VGroup(pixel, text))
        return kernel
    
    def create_feature_maps(self):
        feature_maps = VGroup()
        for _ in range(4):
            fm = VGroup()
            for i in range(4):
                for j in range(4):
                    pixel = Square(side_length=0.4, fill_color=GREEN,
                                 fill_opacity=np.random.uniform(0.3, 0.8),
                                 stroke_color=WHITE, stroke_width=1)
                    pixel.move_to([j*0.4 - 0.6, -i*0.4 + 0.6, 0])
                    fm.add(pixel)
            feature_maps.add(fm)
        feature_maps.arrange(RIGHT, buff=0.3)
        return feature_maps
    
    def create_fc_layers(self):
        layers = VGroup()
        
        # 输入层
        input_layer = VGroup()
        for i in range(5):
            circle = Circle(radius=0.15, fill_color=BLUE, fill_opacity=0.7,
                          stroke_color=WHITE, stroke_width=1)
            circle.move_to([0, -i*0.4 + 0.8, 0])
            input_layer.add(circle)
        layers.add(input_layer)
        
        # 隐藏层
        hidden_layer = VGroup()
        for i in range(3):
            circle = Circle(radius=0.15, fill_color=GREEN, fill_opacity=0.7,
                          stroke_color=WHITE, stroke_width=1)
            circle.move_to([1.5, -i*0.4 + 0.4, 0])
            hidden_layer.add(circle)
        layers.add(hidden_layer)
        
        # 输出层
        output_layer = VGroup()
        for i in range(3):
            circle = Circle(radius=0.15, fill_color=RED, fill_opacity=0.7,
                          stroke_color=WHITE, stroke_width=1)
            circle.move_to([3.0, -i*0.4 + 0.4, 0])
            output_layer.add(circle)
        layers.add(output_layer)
        
        # 连接线
        for inp in input_layer:
            for hid in hidden_layer:
                line = Line(inp.get_center(), hid.get_center(),
                          stroke_width=1.5, color=WHITE)
                layers.add(line)
        
        for hid in hidden_layer:
            for out in output_layer:
                line = Line(hid.get_center(), out.get_center(),
                          stroke_width=1.5, color=WHITE)
                layers.add(line)
        
        return layers

if __name__ == "__main__":
    scene = ImprovedImageRecognition()
    scene.render()